#!/usr/bin/env python
# -*- coding: UTF-8 -*-
#
# Copyright [Gtlions Lai].
# Create Date:
# Update Date:
"""summarization ahout this script.

detail ahout this script

    Class(): summarization about Class
    ...
    function(): summarization about function
    ...
"""
__authors__ = '"Gtlions Lai" <gtlions.l@qq.com>'

import time
import sys
import platform

if platform.system() == 'Windows':
    from psycopg2 import *
elif platform.system() == 'Linux':
    from pygresql.pgdb import *
import datetime
import dsn


class PartitionManager():
    def __init__(self):
        self.db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
        self.cur = self.db.cursor()

    def partition_add_months_integer(relation, reserve):
        relation = str(relation)
        reserve = int(reserve)
        self.cur.execute("select to_char(current_date,'yyyymm')")
        partitionrangestart = self.cur.fetchone()[0]
        self.cur.execute("select to_char(current_date+interval '1 months','yyyymm')")
        partitionrangeend = self.cur.fetchone()[0]
        self.cur.execute("select count(*) from pg_tables where schemaname||'.'||tablename=lower('" + relation + "')")
        v_cnt = self.cur.fetchone()[0]
        if v_cnt:
            pass
        else:
            print 'ERROR: relation "' + relation + '" does not exist'
            return
        if reserve > 0:
            pass
        else:
            print 'ERROR: parameter "i_partition_add_cnt" must greater than 0,but you input ' + str(reserve)
            return

        self.cur.execute(
            "select pk_name_m,pk_name_m_suffix,p_m_parameter from gpmg.manager_table where schemaname||'.'||tablename=lower('" + relation + "')")
        ret = self.cur.fetchone()
        pk_name_m = ret[0]
        pk_name_m_suffix = ret[1]
        p_m_parameter = ret[2]

        sql_exists_partitions = "select 'alter table " + relation + " add partition '|| partitionname ||' start ('||regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')||') end ('||regexp_replace(partitionrangeend,'(::|numeric|bigint)+','','g')||')" + p_m_parameter + ";' from pg_partitions where partitionschemaname||'.'||tablename='" + relation + "' and regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')::int>=" + partitionrangestart + " and regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')::int<299901"
        self.cur.execute("select to_char(current_date,'" + pk_name_m_suffix + "')")
        p_name_m = pk_name_m + self.cur.fetchone()[0]
        sql_next_partitions = ""
        for i in range(1, reserve + 1):
            sql_next_partitions += "select 'alter table " + relation + " add partition " + p_name_m + " start (" + partitionrangestart + ") end (" + partitionrangeend + ")" + p_m_parameter + ";'"
            if i < reserve:
                sql_next_partitions += " union all "

            self.cur.execute(
                "select to_char(current_date+interval '" + str(i) + " months' ,'" + pk_name_m_suffix + "')")
            pk_name_m_suffix_1 = self.cur.fetchone()[0]
            p_name_m = pk_name_m + pk_name_m_suffix_1
            self.cur.execute(
                "select to_char(current_date+interval '" + str(i) + " months' ,'" + pk_name_m_suffix + "')")
            partitionrangestart = self.cur.fetchone()[0]
            self.cur.execute(
                "select to_char(current_date+interval '" + str(i + 1) + " months' ,'" + pk_name_m_suffix + "')")
            partitionrangeend = self.cur.fetchone()[0]

        self.cur.execute("select nextval('gpmg.seq_manager_part_mid')")
        manager_part_mid_sid = self.cur.fetchone()[0]
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t1) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_exists_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t2) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_next_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t3) select sid,t2 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + " except select sid,t1 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + ";")
        self.cur.execute(
            "select t3 from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid) + " and t3 is not null")
        v_msg = "Done.Nothing."
        for record in self.cur.fetchall():
            sql_add_partitions = record[0]
            self.cur.execute(sql_add_partitions)
            v_msg = v_msg.replace('Nothing.', '')
        self.cur.execute("delete from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid))
        self.cur.close()
        self.db.commit()
        self.db.close()
        return v_msg

    def partition_add_day_date(relation, reserve):
        relation = str(relation)
        reserve = int(reserve)
        db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
        cur = self.db.cursor()
        self.cur.execute("select to_char(current_date,'yyyy-mm-dd')")
        partitionrangestart = self.cur.fetchone()[0]
        self.cur.execute("select to_char(current_date+1,'yyyy-mm-dd')")
        partitionrangeend = self.cur.fetchone()[0]
        self.cur.execute("select count(*) from pg_tables where schemaname||'.'||tablename=lower('" + relation + "')")
        v_cnt = self.cur.fetchone()[0]
        if v_cnt:
            pass
        else:
            print 'ERROR: relation "' + relation + '" does not exist'
            return
        if reserve > 0:
            pass
        else:
            print 'ERROR: parameter "i_partition_add_cnt" must greater than 0,but you input ' + str(reserve)
            return

        self.cur.execute(
            "select pk_name_d,pk_name_d_suffix,p_d_parameter from gpmg.manager_table where schemaname||'.'||tablename=lower('" + relation + "')")
        ret = self.cur.fetchone()
        pk_name_d = ret[0]
        pk_name_d_suffix = ret[1]
        p_d_parameter = ret[2]
        sql_exists_partitions = "select 'alter table " + relation + " add partition '|| partitionname ||' start ('||partitionrangestart||') end ('||partitionrangeend||') " + p_d_parameter + ";' from pg_partitions where partitionschemaname||'.'||tablename='" + relation + "' and to_date(substr(partitionrangestart,2,11),'yyyy-mm-dd')>='" + partitionrangestart + "'::date"
        self.cur.execute("select to_char(current_date,'" + pk_name_d_suffix + "')")
        p_name_d = pk_name_d + self.cur.fetchone()[0]
        sql_next_partitions = ""
        for i in range(1, reserve + 1):
            sql_next_partitions += "select 'alter table " + relation + " add partition " + p_name_d + " start (''" + partitionrangestart + "''::date) end (''" + partitionrangeend + "''::date) " + p_d_parameter + ";'"
            if i < reserve:
                sql_next_partitions += " union all "

            self.cur.execute("select to_char(current_date+" + str(i) + ",'" + pk_name_d_suffix + "')")
            pk_name_d_suffix_1 = self.cur.fetchone()[0]
            p_name_d = pk_name_d + pk_name_d_suffix_1
            self.cur.execute("select to_char(current_date+" + str(i) + ",'yyyy-mm-dd')")
            partitionrangestart = self.cur.fetchone()[0]
            self.cur.execute("select to_char(current_date+ " + str(i + 1) + ",'yyyy-mm-dd')")
            partitionrangeend = self.cur.fetchone()[0]
        self.cur.execute("select nextval('gpmg.seq_manager_part_mid')")
        manager_part_mid_sid = self.cur.fetchone()[0]
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t1) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_exists_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t2) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_next_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t3) select sid,t2 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + " except select sid,t1 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + ";")
        self.cur.execute(
            "select t3 from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid) + " and t3 is not null")
        v_msg = "Done.Nothing."
        for record in self.cur.fetchall():
            sql_add_partitions = record[0]
            self.cur.execute(sql_add_partitions)
            v_msg = v_msg.replace('Nothing.', '')
        self.cur.execute("delete from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid))
        self.cur.close()
        self.db.commit()
        self.db.close()
        return v_msg

    def partition_add_day_integer(relation, reserve):
        relation = str(relation)
        reserve = int(reserve)
        db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
        cur = self.db.cursor()
        self.cur.execute("select to_char(current_date,'yyyymmdd')")
        partitionrangestart = self.cur.fetchone()[0]
        self.cur.execute("select to_char(current_date+1,'yyyymmdd')")
        partitionrangeend = self.cur.fetchone()[0]
        self.cur.execute("select count(*) from pg_tables where schemaname||'.'||tablename=lower('" + relation + "')")
        v_cnt = self.cur.fetchone()[0]
        if v_cnt:
            pass
        else:
            print 'ERROR: relation "' + relation + '" does not exist'
            return
        if reserve > 0:
            pass
        else:
            print 'ERROR: parameter "i_partition_add_cnt" must greater than 0,but you input ' + str(reserve)
            return

        self.cur.execute(
            "select pk_name_d,pk_name_d_suffix,p_d_parameter from gpmg.manager_table where schemaname||'.'||tablename=lower('" + relation + "')")
        ret = self.cur.fetchone()
        pk_name_d = ret[0]
        pk_name_d_suffix = ret[1]
        p_d_parameter = ret[2]
        sql_exists_partitions = "select 'alter table " + relation + " add partition '|| partitionname ||' start ('||regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')||') end ('||regexp_replace(partitionrangeend,'(::|numeric|bigint)+','','g')||') '||'" + p_d_parameter + "'||';' from pg_partitions where partitionschemaname||'.'||tablename='" + relation + "' and regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')::int>=" + partitionrangestart + " and regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')::int<29990101"
        self.cur.execute("select to_char(current_date,'" + pk_name_d_suffix + "')")
        p_name_d = pk_name_d + self.cur.fetchone()[0]
        sql_next_partitions = ""
        for i in range(1, reserve + 1):
            sql_next_partitions += "select 'alter table " + relation + " add partition " + p_name_d + " start (" + partitionrangestart + ") end (" + partitionrangeend + ") " + p_d_parameter + ";'"
            if i < reserve:
                sql_next_partitions += " union all "

            self.cur.execute("select to_char(current_date+" + str(i) + ",'" + pk_name_d_suffix + "')")
            pk_name_d_suffix_1 = self.cur.fetchone()[0]
            p_name_d = pk_name_d + pk_name_d_suffix_1
            self.cur.execute("select to_char(current_date+" + str(i) + ",'yyyymmdd')")
            partitionrangestart = self.cur.fetchone()[0]
            self.cur.execute("select to_char(current_date+ " + str(i + 1) + ",'yyyymmdd')")
            partitionrangeend = self.cur.fetchone()[0]
        self.cur.execute("select nextval('gpmg.seq_manager_part_mid')")
        manager_part_mid_sid = self.cur.fetchone()[0]
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t1) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_exists_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t2) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_next_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t3) select sid,t2 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + " except select sid,t1 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + ";")
        self.cur.execute(
            "select t3 from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid) + " and t3 is not null")
        v_msg = "Done.Nothing."
        for record in self.cur.fetchall():
            sql_add_partitions = record[0]
            self.cur.execute(sql_add_partitions)
            v_msg = v_msg.replace('Nothing.', '')
        self.cur.execute("delete from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid))
        self.cur.close()
        self.db.commit()
        self.db.close()
        return v_msg

    def partition_add_hour_integer(relation, reserve):
        relation = str(relation)
        reserve = int(reserve)
        db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
        cur = self.db.cursor()
        self.cur.execute("select to_char(current_date,'yyyymmddhh24')")
        partitionrangestart = self.cur.fetchone()[0]
        self.cur.execute("select to_char(current_date+1,'yyyymmddhh24')")
        partitionrangeend = self.cur.fetchone()[0]
        self.cur.execute("select count(*) from pg_tables where schemaname||'.'||tablename=lower('" + relation + "')")
        v_cnt = self.cur.fetchone()[0]
        if v_cnt:
            pass
        else:
            print 'ERROR: relation "' + relation + '" does not exist'
            return
        if reserve > 0:
            pass
        else:
            print 'ERROR: parameter "i_partition_add_cnt" must greater than 0,but you input ' + str(reserve)
            return

        self.cur.execute(
            "select pk_name_h,pk_name_h_suffix,p_h_parameter from gpmg.manager_table where schemaname||'.'||tablename=lower('" + relation + "')")
        ret = self.cur.fetchone()
        pk_name_h = ret[0]
        pk_name_h_suffix = ret[1]
        p_h_parameter = ret[2]
        start_date = datetime.date.today()
        end_date = datetime.date.today() + datetime.timedelta(days=reserve)
        sql_exists_partitions = "select 'alter table " + relation + " add partition '|| partitionname||' start ('||regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')||') end ('||regexp_replace(partitionrangeend,'(::|numeric|bigint)+','','g')||') '||'" + p_h_parameter + "'||';' from pg_partitions where partitionschemaname||'.'||tablename='" + relation + "' and regexp_replace(partitionrangestart,'(::|numeric|bigint)+','','g')::int>=" + partitionrangestart
        sql_next_partitions = "select 'alter table " + relation + " add partition " + pk_name_h + "'||a.rq||b.h||' start ('||a.rq||b.h||') end ('||(case when b.h='23' then a.rqn else a.rq end)||b.hn||') " + p_h_parameter + ";' from gpmg.manager_date a,gpmg.manager_hour b where a.rqd>='" + str(
            start_date) + "' and a.rqd<'" + str(end_date) + "' order by a.id,b.h"
        self.cur.execute("select nextval('gpmg.seq_manager_part_mid')")
        manager_part_mid_sid = self.cur.fetchone()[0]
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t1) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_exists_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t2) select " + str(
                manager_part_mid_sid) + ",* from (" + sql_next_partitions + ") a ;")
        self.cur.execute(
            "insert into gpmg.manager_part_mid(sid,t3) select sid,t2 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + " except select sid,t1 from gpmg.manager_part_mid where sid=" + str(
                manager_part_mid_sid) + ";")
        self.cur.execute(
            "select t3 from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid) + " and t3 is not null")
        v_msg = "Done.Nothing."
        for record in self.cur.fetchall():
            sql_add_partitions = record[0]
            self.cur.execute(sql_add_partitions)
            v_msg = v_msg.replace('Nothing.', '')
        self.cur.execute("delete from gpmg.manager_part_mid where sid=" + str(manager_part_mid_sid))
        self.cur.close()
        self.db.commit()
        self.db.close()
        return v_msg

    def partition_drop():
        db = connect(host=dsn.host, database=dsn.database, user=dsn.user, password=dsn.password)
        cur = self.db.cursor()
        self.cur.execute("select  to_char(now(),'yyyymmddhh24miss');")
        log_id = self.cur.fetchone()[0]
        sql = "select schemaname,tablename,partlevel,p_type,pk_name_m,pk_name_d, pk_name_h, pk_name_m_suffix, pk_name_d_suffix, pk_name_h_suffix, retention_m, retention_d, retention_h from gpmg.manager_table where debug=0 and (status=1 or (status=0 and expire>current_date) ) and (oper_type='del' or oper_type='all') order by id"
        self.cur.execute(sql)
        serl = 0
        if len(sys.argv) == 2:
            ds = int(sys.argv[1])
        else:
            ds = 0
        v_delta = datetime.timedelta(days=ds)
        v_current_date = (datetime.date.today() + v_delta).strftime('%Y-%m-%d')

        for ret in self.cur.fetchall():
            v_schemaname = ret[0]
            v_tablename = ret[1]
            relation = v_schemaname + "." + v_tablename
            sql_insert_log = "insert into gpmg.manager_table_operlog select '" + v_schemaname + "','" + v_tablename + "', now(),'d','" + log_id + "';"
            sql_update_log = "update gpmg.manager_table set drop_result=now() where schemaname||'.'||tablename='" + relation + "';"
            serl += 1
            print "|" + str(datetime.datetime.now()).center(25) + "|" + str(serl).center(5) + "|" + relation.center(
                45) + "|"
            v_partlevel = ret[2]
            v_p_type = ret[3]
            v_pk_name_m = ret[4]
            v_pk_name_d = ret[5]
            v_pk_name_h = ret[6]
            v_pk_name_m_suffix = ret[7]
            v_pk_name_d_suffix = ret[8]
            v_pk_name_h_suffix = ret[9]
            v_retention_m = ret[10]
            v_retention_d = ret[11]
            v_retention_h = ret[12]
            sql2 = "select count(*) from pg_tables where schemaname||'.'||tablename = '" + relation + "'"
            self.cur.execute(sql2)
            is_exists = self.cur.fetchone()[0]
            if is_exists == 0:
                print "Not exists this table!\n"
                continue
            sql2 = "select count(1) from pg_locks where mode in ('AccessExclusiveLock','ExclusiveLock','RowExclusiveLock') and relation = '" + relation + "'::regclass and gp_segment_id='-1' and granted=True;"
            is_lock = 0
            for i in range(5):
                self.cur.execute(sql2)
                is_lock = self.cur.fetchone()[0]
                if is_lock == 0:
                    if i > 0:
                        print ''
                    break
                print " sleep 10's ",
                time.sleep(10)
            if is_lock > 0:
                print "Can't lock this table!"
                continue
            if (',' + v_p_type + ',').find(',m,') >= 0:
                v_retention = str(v_retention_m)
                v_pk_name = v_pk_name_m
                v_pk_name_suffix = v_pk_name_m_suffix
                sql = "select '" + v_pk_name + "'||to_char('" + v_current_date + "'::date-interval '" + v_retention + " months' ,'" + v_pk_name_suffix + "')";
                self.cur.execute(sql)
                partitionname = self.cur.fetchone()[0]
                sql = "select count(*) from pg_partitions where schemaname||'.'||tablename='" + relation + "' and partitionname='" + partitionname + "'";
                self.cur.execute(sql)
                is_exists = self.cur.fetchone()[0]
                if is_exists:
                    sql = "select schemaname||'.'||partitiontablename from pg_partitions where schemaname||'.'||tablename='" + relation + "' and partitionname='" + partitionname + "'";
                    partitiontablename = self.cur.fetchone()[0]
                    self.cur.execute("select pg_total_relation_size('" + partitiontablename + "')")
                    retsize = self.cur.fetchone()[0]
                    self.cur.execute("select nextval('gpmg.seq_unload_size')")
                    retid = self.cur.fetchone()[0]
                    self.cur.execute(
                        "insert into gpmg.unload_size(id,rq,tablename,operation,tablesize) values(%s,%s,%s,'drop table',%s)",
                        (retid, datetime.date.today(), partitiontablename, retsize))
                    self.cur.execute("alter table " + relation + " drop partition " + partitionname)
                    self.cur.execute(sql_update_log)
                    self.cur.execute(sql_insert_log)
            if (',' + v_p_type + ',').find(',d,') >= 0:
                v_retention = str(v_retention_d)
                v_pk_name = v_pk_name_d
                v_pk_name_suffix = v_pk_name_d_suffix
                sql = "select '" + v_pk_name + "'||to_char('" + v_current_date + "'::date-" + v_retention + ",'" + v_pk_name_suffix + "')";
                self.cur.execute(sql)
                partitionname = self.cur.fetchone()[0]
                sql = "select count(*) from pg_partitions where schemaname||'.'||tablename='" + relation + "' and partitionname='" + partitionname + "'";
                self.cur.execute(sql)
                is_exists = self.cur.fetchone()[0]
                if is_exists:
                    sql = "select schemaname||'.'||partitiontablename from pg_partitions where schemaname||'.'||tablename='" + relation + "' and partitionname='" + partitionname + "'";
                    self.cur.execute(sql)
                    partitiontablename = self.cur.fetchone()[0]
                    self.cur.execute("select pg_total_relation_size('" + partitiontablename + "')")
                    retsize = self.cur.fetchone()[0]
                    self.cur.execute("select nextval('gpmg.seq_unload_size')")
                    retid = self.cur.fetchone()[0]
                    self.cur.execute(
                        "insert into gpmg.unload_size(id,rq,tablename,operation,tablesize) values(%s,%s,%s,'drop table',%s)",
                        (retid, datetime.date.today(), partitiontablename, retsize))
                    self.cur.execute("alter table " + relation + " drop partition " + partitionname)
                    self.cur.execute(sql_update_log)
                    self.cur.execute(sql_insert_log)
            if v_p_type == 'h':
                v_retention = str(v_retention_h)
                v_pk_name = v_pk_name_h
                v_pk_name_suffix = v_pk_name_h_suffix
                for i in range(24):
                    sql = "select '" + v_pk_name + "'||to_char('" + v_current_date + "'::date -" + v_retention + " +interval '" + str(
                        i) + " hours' ,'yyyymmddhh24')";
                    self.cur.execute(sql)
                    partitionname = self.cur.fetchone()[0]
                    sql = "select count(*) from pg_partitions where schemaname||'.'||tablename='" + relation + "' and partitionname='" + partitionname + "'";
                    self.cur.execute(sql)
                    is_exists = self.cur.fetchone()[0]
                    if is_exists:
                        sql = "select schemaname||'.'||partitiontablename from pg_partitions where schemaname||'.'||tablename='" + relation + "' and partitionname='" + partitionname + "'";
                        self.cur.execute(sql)
                        partitiontablename = self.cur.fetchone()[0]
                        self.cur.execute("select pg_total_relation_size('" + partitiontablename + "')")
                        retsize = self.cur.fetchone()[0]
                        self.cur.execute("select nextval('gpmg.seq_unload_size')")
                        retid = self.cur.fetchone()[0]
                        self.cur.execute(
                            "insert into gpmg.unload_size(id,rq,tablename,operation,tablesize) values(%s,%s,%s,'drop table',%s)",
                            (retid, datetime.date.today(), partitiontablename, retsize))
                        self.cur.execute("alter table " + relation + " drop partition " + partitionname)
                self.cur.execute(sql_update_log)
                self.cur.execute(sql_insert_log)
            self.db.commit()
        self.db.commit()
        self.db.close()


if __name__ == "__main__":
    pass
